

from torchvision import transforms

import numpy as np
import librosa
import glob

import matplotlib.pyplot as plt
# from matplotlib.cm import get_cmap



viridis_cmap = plt.get_cmap('viridis')
color_map = viridis_cmap.colors
color_map = (np.array(color_map) * 255).astype(np.uint8)
s_path = "D:/Mind+/ext_work/ext_librosa_test/example/打开"
time=2
X =[]
y =[]
 # Load the file
audio_files = glob.glob(s_path + "/*/*.m4a") + glob.glob(s_path + "/*/*.mp3") + glob.glob(s_path + "/*/*.wav") + glob.glob(s_path + "/*/*.flac")
print(audio_files)
# read file
for file in audio_files:
    print(file)
    # 获得标签
    y.append(file.split('/')[-2])
    # 读取音频
    wave,fs = librosa.load(file, sr = 16000)
    # time对应的样点数
    sample = int(time * fs)
    # 若音频短于time的样点数，则在音频后面补零，若音频长于time的样点数，则对音频截断
    if wave.size <= sample:
        wave = np.concatenate((wave,np.array((sample - wave.size) * [0])))
    else:
        wave = wave[0:sample]
    #打印音频长度
    print(wave.shape)
    # 获取梅尔频谱
    spc = librosa.feature.melspectrogram(wave, sr=fs, n_fft=512)
    #转化为log形式
    spc = librosa.power_to_db(spc, ref=np.max)
    #数据归一化
    spec_new = (((spc+80)/80)*255).astype(np.uint8)
    h,w = spec_new.shape
    rgb_matrix = np.array([color_map[i] for i in spec_new.flatten()]).reshape(h, w, 3)
    spec_rgb = rgb_matrix/255
    # 添加到数据
    X.append(spec_rgb)


out_path = "D:/Mind+/ext_work/ext_librosa_test/example/频谱图"
for i,j in enumerate(y):
    path = out_path+"/" + str(j) + "/"
    test = X[i]
    test = transforms.ToTensor()(test)
    test = transforms.Resize((224, 224))(test)
    plt.imsave(path+str(i)+'.jpg', np.array(test).transpose(1,2,0))
print('验证数据保存成功')
